#include "ipu_image_preprocessor.hpp"
#include <iostream>
#include <stdexcept>
#include <iomanip>

namespace protonet {

IPUImagePreprocessor::IPUImagePreprocessor(const std::pair<int, int>& target_size,
                                          const PreprocessConfig* config)
    : target_size_(target_size) {
    
    if (config != nullptr) {
        config_ = *config;
    }
    
    // 更新配置中的尺寸信息
    config_.resize_width = target_size_.first;
    config_.resize_height = target_size_.second;
    
    std::cout << "🖼️  IPU图像预处理器初始化完成" << std::endl;
    std::cout << "📐 目标尺寸: " << target_size_.first << "x" << target_size_.second << std::endl;
    std::cout << "🎨 输入格式: " << config_.input_format << std::endl;
    std::cout << "🔧 归一化: " << (config_.do_normalize ? "启用" : "禁用") << std::endl;
    if (config_.do_normalize) {
        std::cout << "📊 均值: [" << std::fixed << std::setprecision(3) 
                  << config_.mean_values[0] << ", " << config_.mean_values[1] << ", " 
                  << config_.mean_values[2] << "]" << std::endl;
        std::cout << "📊 标准差: [" << std::fixed << std::setprecision(3) 
                  << config_.std_values[0] << ", " << config_.std_values[1] << ", " 
                  << config_.std_values[2] << "]" << std::endl;
    }
}

void IPUImagePreprocessor::setTargetSize(const std::pair<int, int>& target_size) {
    target_size_ = target_size;
    config_.resize_width = target_size.first;
    config_.resize_height = target_size.second;
}

std::vector<float> IPUImagePreprocessor::preprocessImage(const std::string& image_path) {
    // 读取图像
    cv::Mat image = cv::imread(image_path);
    if (image.empty()) {
        throw std::runtime_error("无法读取图像文件: " + image_path);
    }
    
    return processImage(image);
}

std::vector<float> IPUImagePreprocessor::preprocessImage(const cv::Mat& image) {
    if (image.empty()) {
        throw std::runtime_error("输入图像为空");
    }
    
    return processImage(image.clone());
}

std::vector<float> IPUImagePreprocessor::processImage(cv::Mat image) {
    try {
        // 1. 处理灰度图像转为3通道
        if (image.channels() == 1) {
            cv::cvtColor(image, image, cv::COLOR_GRAY2BGR);
        } else if (image.channels() == 4) {
            cv::cvtColor(image, image, cv::COLOR_BGRA2BGR);
        }
        
        // 2. 调整图像尺寸
        if (image.rows != config_.resize_height || image.cols != config_.resize_width) {
            cv::resize(image, image, cv::Size(config_.resize_width, config_.resize_height), 
                      0, 0, cv::INTER_LINEAR);
        }
        
        // 3. 颜色格式转换
        image = convertColorFormat(image);
        
        // 4. 归一化处理
        cv::Mat normalized_image = applyNormalization(image);
        
        // 5. 转换为float向量
        std::vector<float> data;
        int total_elements = normalized_image.rows * normalized_image.cols * normalized_image.channels();
        data.reserve(total_elements);
        
        // 按NHWC格式存储数据（N=1, HWC）
        if (normalized_image.isContinuous()) {
            float* img_data = (float*)normalized_image.data;
            data.assign(img_data, img_data + total_elements);
        } else {
            // 处理非连续内存的情况
            for (int h = 0; h < normalized_image.rows; ++h) {
                for (int w = 0; w < normalized_image.cols; ++w) {
                    cv::Vec3f pixel = normalized_image.at<cv::Vec3f>(h, w);
                    data.push_back(pixel[0]); // B or R
                    data.push_back(pixel[1]); // G
                    data.push_back(pixel[2]); // R or B
                }
            }
        }
        
        // 6. 如果需要，转换为NCHW格式
        if (config_.to_nchw) {
            data = convertToNCHW(data, config_.resize_height, config_.resize_width, 3);
        }
        
        return data;
        
    } catch (const cv::Exception& e) {
        throw std::runtime_error("OpenCV错误: " + std::string(e.what()));
    } catch (const std::exception& e) {
        throw std::runtime_error("图像预处理错误: " + std::string(e.what()));
    }
}

cv::Mat IPUImagePreprocessor::applyNormalization(const cv::Mat& image) {
    if (!config_.do_normalize) {
        // 如果不需要归一化，转换为float类型并返回
        cv::Mat float_image;
        image.convertTo(float_image, CV_32F);
        return float_image;
    }
    
    cv::Mat float_image, normalized_image;
    image.convertTo(float_image, CV_32F);
    
    if (config_.norm_scheme == "subtract_first") {
        // 方案1: (pixel - mean) / std
        cv::Scalar mean_scalar(config_.mean_values[0], config_.mean_values[1], config_.mean_values[2]);
        cv::Scalar std_scalar(config_.std_values[0], config_.std_values[1], config_.std_values[2]);
        
        normalized_image = (float_image - mean_scalar);
        
        // 如果标准差不全为1，则进行除法操作
        if (config_.std_values[0] != 1.0f || config_.std_values[1] != 1.0f || config_.std_values[2] != 1.0f) {
            std::vector<cv::Mat> channels;
            cv::split(normalized_image, channels);
            
            for (int i = 0; i < 3; ++i) {
                channels[i] = channels[i] / config_.std_values[i];
            }
            
            cv::merge(channels, normalized_image);
        }
    } else {
        // 默认归一化方案
        cv::Scalar mean_scalar(config_.mean_values[0], config_.mean_values[1], config_.mean_values[2]);
        normalized_image = float_image - mean_scalar;
    }
    
    return normalized_image;
}

cv::Mat IPUImagePreprocessor::convertColorFormat(const cv::Mat& image) {
    cv::Mat converted_image = image;
    
    // 根据配置进行颜色格式转换
    if (config_.to_rgb && config_.input_format == "BGR") {
        cv::cvtColor(converted_image, converted_image, cv::COLOR_BGR2RGB);
    } else if (!config_.to_rgb && config_.input_format == "RGB") {
        cv::cvtColor(converted_image, converted_image, cv::COLOR_RGB2BGR);
    }
    
    return converted_image;
}

std::vector<float> IPUImagePreprocessor::convertToNCHW(const std::vector<float>& data, 
                                                      int height, int width, int channels) {
    std::vector<float> nchw_data(data.size());
    
    // 从HWC转换为CHW
    for (int c = 0; c < channels; ++c) {
        for (int h = 0; h < height; ++h) {
            for (int w = 0; w < width; ++w) {
                int hwc_index = h * width * channels + w * channels + c;
                int chw_index = c * height * width + h * width + w;
                nchw_data[chw_index] = data[hwc_index];
            }
        }
    }
    
    return nchw_data;
}

} // namespace protonet